import itertools
import logging
import re
from typing import List, Mapping

import networkx as nx

from chernc.analyzer.rust_simple_ast import parse_on_treesit_nodes
from chernc.analyzer.semantic_piece import *
from chernc.analyzer.source_based_piece import *
from chernc.analyzer.source_based_piece import TypedefStructurePiece
from chernc.analyzer.target_based_piece import *
from chernc.constants import DEFAULT_LOGGER_DIR, DEFAULT_LOGGER_FILE_NAME
from chernc.core.context_manager import SnippetResult
from chernc.core.result import Marks, Match
from chernc.logging import build_logger
from chernc.parser.rust_parser import parse_rust
from chernc.slicer.code_chunk import CodeChunk

logger: logging.Logger = build_logger(
    __name__ + ".c_rust_matcher", logger_filename=DEFAULT_LOGGER_FILE_NAME, logger_dir=DEFAULT_LOGGER_DIR
)


# input: list of lists of slices
# output: components
def get_clusters(relations: Iterable[Iterable[CodeSlice]]) -> List[List[CodeSlice]]:
    # Create an undirected graph
    G = nx.Graph()

    # Add edges for each list
    for relation_list in relations:
        G.add_edges_from(itertools.combinations_with_replacement(relation_list, 2))
    return list(nx.connected_components(G))


# TODO: put these code into CherncMatcher._match, and use logger instead
# These infomation are for debugging


def explain_match(match: Match) -> str:
    """
    Explain how the source chunk matches with this chunk result.
    """

    matched_names = []
    for r in match.values():
        matched_names.extend(r.names)
    return (
        "\n".join(f"===piece text===\n{t.text}\n====matched text===\n{mt.text}\n===piece end===\n" for t, mt in match.items())
        + f"\n==================names==============\n"
        + "\n".join(matched_names)
    )


class MatchFail(Exception):
    def __init__(self, msg):
        super().__init__(f"Match Fail: {msg}")


def text_of_match(match: Match) -> str:
    results = sorted(set(match.values()), key=lambda r: min(s.level for s in r.slices))
    return "\n\n".join(r.text for r in results)


class Matcher(ABC):
    @abstractmethod
    def match(self, result: str) -> Match:
        """
        This method gurantee each slices in the chunk is matched with not DummySliceResult.
        Raises:
            MatchFail: if the matching fails
        """
        raise NotImplementedError("The match method in Matcher is abstract")

    @abstractmethod
    def try_to_match(self, result: str) -> Match:
        """
        This method may return the chunk result with DummySnippetResult. No MatchFail will be raised.
        """
        raise NotImplementedError("The try_to_match method in Matcher is abstract")


def raiseMergeFail(msg):
    raise MatchFail(msg)


# the matcher
# logically, C code -> matcher <- rust code
class CherncMatcher(Matcher):
    pieces: List[SemanticPiece]
    slices: Set[CodeSlice]

    @staticmethod
    def createforChunk(chunk: CodeChunk) -> "CherncMatcher":
        """
        Returns:
            CherncMatcher: if the chunk is valid
            None: if the chunk is invalid
        """
        return CherncMatcher(chunk.slices)

    def __init__(self, slices: Iterable[CodeSlice]) -> None:
        """
        Returns:
            CherncMatcher: if the chunk is valid
            None: if the chunk is invalid
        """
        slices = set(slices)
        semanticpieces: List[SemanticPiece] = []

        for s in slices:
            pieces = extractPiecesFromSlice(s)
            semanticpieces += pieces
        self.pieces = semanticpieces
        self.slices = slices

    def _match(self, result: str, match_fail: Callable[[str], None]) -> Match:
        rust_treenode: RustNode = parse_rust(result).root_node  # type: ignore
        rust_nodes: Set[RustNode] = set(parse_on_treesit_nodes(rust_treenode.children))  # type: ignore
        # self._rust_nodes = parse_on_treesit_nodes(rust_treenode.children)

        pieces = set(self.pieces)
        matches: List[Tuple[Set[CodeSlice], MatchData]] = []
        # first, accumulate sourcebased pieces
        tds: List[TypedefStructurePiece] = TypedefStructurePiece.extract_from_pieces(pieces)
        for t in tds:
            if (ok := t.match(rust_nodes)) is not None:
                matches.append((set(t.get_slices()), ok))
                rust_nodes.difference_update(ok.rust_nodes)
                pieces.remove(t.alias_piece)
                pieces.remove(t.td_piece)
            else:
                match_fail(t.matching_failure_description())
                continue
        # then, match iteratively
        impls = ImplPiece.extract_from_rust_nodes(rust_nodes)

        for i in impls:
            ret, to_remove = ImplPiece.match_piece_for_node(pieces, i)
            if ret is not None:
                matches.append(({p.father for p in to_remove if p.father is not None}, ret))
                pieces.difference_update(to_remove)
            else:
                match_fail(ImplPiece.mk_fail_description(i))
                continue
        # finally, gradually match all pieces, make sure that all

        for p in pieces:
            if not p.father:
                continue
            if (ret := p.match(rust_nodes)) is not None:
                matches.append(({p.father}, ret))
                rust_nodes.difference_update(ret.rust_nodes)
            else:
                match_fail(p.matching_failure_description())
                continue

        result_table = {}
        for cluster in get_clusters(slices for slices, _ in matches):
            cluster = set(cluster)

            match_data = merge_matches(md for slices, md in matches if slices & cluster)
            if match_data is None:
                continue
            snippet = SnippetResult(cluster, match_data.text, match_data.names)
            for slice in cluster:
                result_table[slice] = snippet
        # final check
        for s in self.slices - result_table.keys():
            match_fail(f"the source code contains\n{s.text}\n, but they have no match in the translation result.")
        return result_table

    def match(self, result: str) -> Match:
        return self._match(result, match_fail=raiseMergeFail)

    def try_to_match(self, result: str) -> Match:
        return self._match(result, match_fail=logger.debug)


class MarkedSnippetResult(SnippetResult):
    def remove_mark(self) -> SnippetResult:
        text = "\n".join(self.text.splitlines()[1:-1])
        return SnippetResult(self.slices, text, self.names)


MarkMatch = Mapping[CodeSlice, MarkedSnippetResult]


def remove_comments(code: str) -> str:
    code = re.sub(r"//.*", "", code)

    no_star = r"[^*]*"
    # strings endswith a star and not followed by a slash
    star_end = no_star + r"\*(?!/)"
    many_star_end = f"(" + star_end + r")*"
    code = re.sub(r"/\*" + many_star_end + no_star + r"\*/", "", code)
    return code


class MarkMatcher(Matcher):
    """
    The match result of this matcher contains the markers.
    """

    def __init__(self, marks: Marks):
        self.marks = marks.copy()

    def match(self, result: str) -> MarkMatch:
        return self._match(result, raiseMergeFail)

    def try_to_match(self, result: str) -> MarkMatch:
        return self._match(result, logger.error)

    def _get_names(self, code: str) -> Set[str]:
        # remove commented code
        code = remove_comments(code)
        # remove blocks
        while (_code := re.sub(r"\{[^{}]*\}", "", code)) != code:
            code = _code
        # Regular expressions for top-level declarations and global variables
        patterns = [
            r"\bfn\s+(\w+)\s*\(",
            r"\bstruct\s+(\w+)\b",
            r"\benum\s+(\w+)\b",
            r"\btrait\s+(\w+)\b",
            r"\bconst\s+(\w+):",
            r"\bstatic\s+(?:mut\s+)?(\w+):",
            r"\btype\s+(\w+)\s+=",
        ]
        names = set()
        for pattern in patterns:
            names.update(re.findall(pattern, code))
        return names

    def _get_code_of_mark(self, code: str, head: str, tail: str) -> Optional[str]:
        """
        Returns:
            str: the code block between head and tail, the first line must be head, and the last line must be tail
            None: if not found
        """
        start = code.rfind(head + "\n")
        if start == -1:
            return None
        end = code.find(tail, start)
        if end == -1:
            return None

        return code[start:end] + tail

    def _match(self, result: str, match_fail: Callable[[str], None]) -> MarkMatch:
        def isValid(code: str) -> bool:
            for line in code.splitlines():
                # this line is empty
                if not line.startswith(" ") and not line.startswith("\t"):
                    continue
                return True
            return False

        match = {}
        for slice, (head, tail, optional) in self.marks.items():
            # found header and tail
            # LLM 喜欢干两件事：1.把代码全删了 2. 把代码全注释了
            if (res := self._get_code_of_mark(result, head, tail)) and (names := self._get_names(res)):
                match[slice] = MarkedSnippetResult([slice], res, names)
            # not found
            elif not optional:
                match_fail(f"{head} is missing")
                continue

        return match


if __name__ == "__main__":
    s = CodeSlice(pool=None, file="", start_line=0, end_line=0)
    s.text = """int arraylist_insert(ArrayList *arraylist, unsigned int index,
                     ArrayListValue data)
{
	/* Sanity check the index */

	if (index > arraylist->length) {
		return 0;
	}

	/* Increase the size if necessary */

	if (arraylist->length + 1 > arraylist->_alloced) {
		if (!arraylist_enlarge(arraylist)) {
			return 0;
		}
	}

	/* Move the contents of the array forward from the index
	 * onwards */

	memmove(&arraylist->data[index + 1],
	        &arraylist->data[index],
	        (arraylist->length - index) * sizeof(ArrayListValue));

	/* Insert the new entry at the index */

	arraylist->data[index] = data;
	++arraylist->length;

	return 1;
}

ArrayList *arraylist_new(unsigned int length)
{
	ArrayList *new_arraylist;

	/* If the length is not specified, use a sensible default */

	if (length <= 0) {
		length = 16;
	}

	/* Allocate the new ArrayList and fill in the fields.  There are
	 * initially no entries. */

	new_arraylist = (ArrayList *) malloc(sizeof(ArrayList));

	if (new_arraylist == NULL) {
		return NULL;
	}

	new_arraylist->_alloced = length;
	new_arraylist->length = 0;

	/* Allocate the data array */

	new_arraylist->data = malloc(length * sizeof(ArrayListValue));

	if (new_arraylist->data == NULL) {
		free(new_arraylist);
		return NULL;
	}

	return new_arraylist;
}

void arraylist_free(ArrayList *arraylist)
{
	/* Do not free if a NULL pointer is passed */

	if (arraylist != NULL) {
		free(arraylist->data);
		free(arraylist);
	}
}"""
    matcher = CherncMatcher([s])
    rust_code = """use std::ptr;
use std::alloc::{self, Layout};

type ArrayListValue = *mut ();

struct ArrayList {
    pub data: *mut ArrayListValue,
    pub length: u32,
    pub _alloced: u32,
}

impl ArrayList {
    fn new(length: u32) -> Option<Box<Self>> {
        // If the length is not specified, use a sensible default
        let length = if length <= 0 { 16 } else { length };

        // Allocate the new ArrayList and fill in the fields. There are initially no entries.
        let layout = Layout::new::<ArrayList>();
        let new_arraylist = unsafe { alloc::alloc(layout) as *mut ArrayList };

        if new_arraylist.is_null() {
            return None;
        }

        // Allocate the data array
        let data_layout = Layout::array::<ArrayListValue>(length as usize).unwrap();
        let data = unsafe { alloc::alloc(data_layout) as *mut ArrayListValue };

        if data.is_null() {
            unsafe { alloc::dealloc(new_arraylist as *mut u8, layout); }
            return None;
        }

        unsafe {
            (*new_arraylist).data = data;
            (*new_arraylist)._alloced = length;
            (*new_arraylist).length = 0;
        }

        Some(unsafe { Box::from_raw(new_arraylist) })
    }

    fn free(arraylist: Option<Box<Self>>) {
        if let Some(mut arraylist) = arraylist {
            let layout = Layout::new::<ArrayList>();
            let data_layout = Layout::array::<ArrayListValue>(arraylist._alloced as usize).unwrap();

            unsafe {
                alloc::dealloc(arraylist.data as *mut u8, data_layout);
                alloc::dealloc(Box::into_raw(arraylist) as *mut u8, layout);
            }
        }
    }

    fn insert(&mut self, index: u32, data: ArrayListValue) -> bool {
        // Sanity check the index
        if index > self.length {
            return false;
        }

        // Increase the size if necessary
        if self.length + 1 > self._alloced {
            if !self.enlarge() {
                return false;
            }
        }

        // Move the contents of the array forward from the index onwards
        unsafe {
            ptr::copy(
                self.data.add(index as usize),
                self.data.add(index as usize + 1),
                (self.length - index) as usize,
            );
        }

        // Insert the new entry at the index
        unsafe {
            *self.data.add(index as usize) = data;
        }
        self.length += 1;

        true
    }

    fn enlarge(&mut self) -> bool {
        let newsize = self._alloced * 2;
        let data_layout = Layout::array::<ArrayListValue>(newsize as usize).unwrap();
        let new_data = unsafe { alloc::realloc(self.data as *mut u8, data_layout, data_layout.size()) as *mut ArrayListValue };

        if new_data.is_null() {
            return false;
        }

        self.data = new_data;
        self._alloced = newsize;

        true
    }

    fn prepend(&mut self, data: ArrayListValue) -> bool {
        self.insert(0, data)
    }
}

fn alloc_test_set_limit(alloc_count: i32) {
    unsafe {
        allocation_limit = alloc_count;
    }
}

fn test_arraylist_new_free() {
    unsafe {
        let mut arraylist: Option<Box<ArrayList>> = None;

        arraylist = ArrayList::new(0);
        assert!(arraylist.is_some());
        ArrayList::free(arraylist);

        arraylist = ArrayList::new(10);
        assert!(arraylist.is_some());
        ArrayList::free(arraylist);

        ArrayList::free(None);

        alloc_test_set_limit(0);
        arraylist = ArrayList::new(0);
        assert!(arraylist.is_none());

        alloc_test_set_limit(1);
        arraylist = ArrayList::new(100);
        assert!(arraylist.is_none());
    }
}"""
    m = matcher.try_to_match(rust_code)
    print(explain_match(m))
